import torch
from models.BiFNet.BiFNet import BiFNet


def torch2onnx(
        weight: str = '',
        output_file: str = '',
        in_chans: int = 1,
        img_size: int = 224):
    # model = BiConv(dim=32, depth=[3, 3, 12, 5], in_chans=1, kernel_size=7, patch_size=4,
    #                    num_classes=16, H=224, W=224, p_h=[8, 4, 2, 1], p_w=[8, 4, 2, 1])
    DEVICE = torch.device("cpu")
    model = torch.load(weight)
    model.eval().to(DEVICE)

    input_shape = (1, in_chans, img_size, img_size)
    class_img = torch.randn(input_shape).to(DEVICE)
    defect_img = torch.randn(input_shape).to(DEVICE)
    dummy_inputs = (class_img, defect_img)

    torch.onnx.export(
        model,
        dummy_inputs,
        output_file,
        input_names=['class_images', 'defect_images'],
        output_names=['output'],
        dynamic_axes={
            'class_images': {0: 'batch_size'},
            'defect_images': {0: 'batch_size'},
            'output':{0: 'batch_size'}
        },
        opset_version=11,
        export_params=True
    )



if __name__ == '__main__':
    weight = "model_123_92.304.pth"
    output_file = "BiFNet_48_33125.onnx"
    torch2onnx(weight, output_file)
